import pyodbc
import numpy as np
from sentence_transformers import SentenceTransformer
import faiss
import pickle
import os
from typing import List, Dict, Tuple
import requests
import json
import time


class QwenVLRAGSystem:
    def __init__(self, db_path: str, qwen_api_url: str = "http://localhost:8000",
                 embedding_model: str = 'all-MiniLM-L6-v2'):
        """
        基于Qwen2.5-VL-8B的RAG系统 - 纯文本版本
        """
        self.db_path = db_path
        self.qwen_api_url = qwen_api_url
        self.connection_string = f'DRIVER={{Microsoft Access Driver (*.mdb, *.accdb)}};DBQ={db_path};'

        # 初始化嵌入模型
        self.embedding_model = SentenceTransformer(embedding_model)

        # 向量索引和产品数据
        self.index = None
        self.products = []
        self.product_texts = []

        # 构建或加载RAG索引
        self.build_or_load_rag_index()

        # 测试API连接
        self.test_api_connection()

    def test_api_connection(self):
        """测试API连接"""
        try:
            response = requests.get(f"{self.qwen_api_url}/health", timeout=10)
            if response.status_code == 200:
                health_data = response.json()
                print(f"✅ API连接成功 - 模型加载: {health_data.get('model_loaded', False)}")
                return True
            else:
                print("❌ API连接失败")
                return False
        except Exception as e:
            print(f"❌ API连接异常: {e}")
            return False

    def build_or_load_rag_index(self):
        """构建或加载RAG索引"""
        index_file = "/mnt/d/project/llt-model/qwen_vl_rag_index.pkl"

        if os.path.exists(index_file):
            print("加载现有的RAG索引...")
            with open(index_file, 'rb') as f:
                data = pickle.load(f)
                self.index = data['index']
                self.products = data['products']
                self.product_texts = data['product_texts']
            print(f"索引加载完成: {len(self.products)} 个产品")
        else:
            print("构建新的RAG索引...")
            self.build_rag_index()

    def build_rag_index(self):
        """构建RAG索引"""
        # 获取所有产品数据
        all_products = self.get_all_products()
        print(f"获取到 {len(all_products)} 个产品")

        # 构建产品文本表示
        self.products = all_products
        self.product_texts = self._create_product_texts(all_products)

        # 生成嵌入向量
        print("生成嵌入向量...")
        embeddings = self.embedding_model.encode(self.product_texts, show_progress_bar=True, batch_size=32)

        # 构建FAISS索引
        dimension = embeddings.shape[1]
        self.index = faiss.IndexFlatIP(dimension)

        # 归一化向量
        faiss.normalize_L2(embeddings)
        self.index.add(embeddings.astype('float32'))

        # 保存索引
        with open("qwen_vl_rag_index.pkl", 'wb') as f:
            pickle.dump({
                'index': self.index,
                'products': self.products,
                'product_texts': self.product_texts
            }, f)

        print(f"RAG索引构建完成: {len(self.products)} 个产品")

    def _create_product_texts(self, products: List[Dict]) -> List[str]:
        """创建产品文本表示"""
        texts = []
        for product in products:
            # 组合多个字段作为文本表示
            product_name = product.get('英文品名', '')
            cn_name = product.get('中文品名', '')
            material = product.get('产品材质', '')
            hs_code = product.get('产品海关编码', '')
            base_tax = product.get('基础关税', '')
            add_tax = product.get('附加关税', '')
            ctns = product.get('件数CTNS', '')
            product_count = product.get('产品申报数量', '')
            weight = product.get('净重KGS', '')
            actual_weight = product.get('收货实重', '')
            volume = product.get('方数', '')
            product_price = product.get('产品申报单价', '')
            total_value = product.get('总价', '')
            container_no = product.get('柜号', '')
            file_path = product.get('文件路径', '')

            text_parts = [f"英文品名: {product_name}"]
            text_parts.append(f"中文品名: {cn_name}")
            text_parts.append(f"产品材质: {material}")
            text_parts.append(f"产品海关编码: {hs_code}")
            text_parts.append(f"基础关税: {base_tax}")
            text_parts.append(f"附加关税: {add_tax}")
            text_parts.append(f"件数CTNS: {ctns}")
            text_parts.append(f"产品申报数量: {product_count}")
            text_parts.append(f"净重KGS: {weight}")
            text_parts.append(f"收货实重: {actual_weight}")
            text_parts.append(f"方数: {volume}")
            text_parts.append(f"产品申报单价: {product_price}")
            text_parts.append(f"总价: {total_value}")
            if container_no:
                text_parts.append(f"柜号: {container_no}")
            text_parts.append(f"文件路径: {file_path}")

            texts.append("; ".join(text_parts))

        return texts

    def get_all_products(self) -> List[Dict]:
        """获取所有产品数据"""
        try:
            conn = pyodbc.connect(self.connection_string)
            cursor = conn.cursor()
            cursor.execute("SELECT * FROM 清关资料")
            columns = [column[0] for column in cursor.description]

            results = []
            for row in cursor.fetchall():
                row_dict = dict(zip(columns, row))
                results.append(row_dict)

            conn.close()
            return results
        except Exception as e:
            print(f"数据库查询失败: {e}")
            return []

    def retrieve_similar_products(self, query: str, top_k: int = 10) -> List[Tuple[Dict, float]]:
        """检索相似产品"""
        if self.index is None or len(self.products) == 0:
            return []

        # 生成查询向量
        query_embedding = self.embedding_model.encode([query])
        faiss.normalize_L2(query_embedding)

        # 搜索相似产品
        scores, indices = self.index.search(query_embedding.astype('float32'), top_k)

        results = []
        for score, idx in zip(scores[0], indices[0]):
            if idx < len(self.products) and score > 0.1:
                results.append((self.products[idx], float(score)))

        return results

    def format_products_table(self, products: List[Tuple[Dict, float]]) -> str:
        """将产品列表格式化为文本表格"""
        if not products:
            return "无匹配商品"

        table_lines = []
        table_lines.append("=" * 120)
        table_lines.append(
            f"{'序号':<4} {'英文品名':<30} {'中文品名':<30} {'产品材质':<30} {'产品海关编码':<30} {'基础关税':<8} {'附加关税':<8} {'件数CTNS':<8} {'产品申报数量':<8} {'净重KGS':<8} {'收货实重':<8} {'方数':<8} {'产品申报单价':<8} {'总价':<8} {'柜号':<30} {'文件路径':<30}")
        table_lines.append("-" * 120)

        for i, (product, score) in enumerate(products, 1):
            product_name = product.get('英文品名', '') or "N/A"
            cn_name = product.get('中文品名', '') or "N/A"
            material = product.get('产品材质', '') or "N/A"
            hs_code = product.get('产品海关编码', '') or "N/A"
            base_tax = product.get('基础关税', '') or "N/A"
            add_tax = product.get('附加关税', '') or "N/A"
            ctns = product.get('件数CTNS', '') or "N/A"
            product_count = product.get('产品申报数量', '') or "N/A"
            weight = product.get('净重KGS', '') or "N/A"
            actual_weight = product.get('收货实重', '') or "N/A"
            volume = product.get('方数', '') or "N/A"
            product_price = product.get('产品申报单价', '') or "N/A"
            total_value = product.get('总价', '') or "N/A"
            container_no = product.get('柜号', '') or "N/A"
            file_path = product.get('文件路径', '') or "N/A"
            table_lines.append(
            f"{i:<4} {product_name:<30} {cn_name:<30} {material:<30} {hs_code:<30} {base_tax:<8} {add_tax:<8} {ctns:<8} {product_count:<8} {weight:<8} {actual_weight:<8} {volume:<8} {product_price:<8} {total_value:<8} {container_no:<30} {file_path:<30}"
            )

        table_lines.append("=" * 120)
        return "\n".join(table_lines)

    def call_qwen_vl_with_text(self, query: str, products_table: str) -> str:
        """使用纯文本调用Qwen-VL进行分析"""

        prompt = f"""
你是一个专业的海关商品分类专家。请根据以下商品匹配结果，为用户查询提供专业的匹配建议。

用户查询: "{query}"

匹配结果表格:
{products_table}

请基于以上匹配结果，完成以下分析任务：

1. **匹配度分析**
   - 分析查询商品与匹配商品的相似程度
   - 评估整体匹配质量（优秀/良好/一般/较差）

2. **推荐建议**
   - 推荐最匹配的前3个商品
   - 说明每个推荐商品的匹配理由
   - 对比关键属性（规格、材质、用途等）

3. **专业建议**
   - 海关分类建议（基于HS编码）
   - 潜在的风险或注意事项
   - 是否需要进一步信息确认

4. **总结**
   - 给出最终的匹配结论
   - 提供后续操作建议

请用专业、清晰的语言回答，结构清晰，重点突出。
"""

        try:
            return self._call_chat_completions(prompt)
        except Exception as e:
            print(f"Qwen-VL分析失败: {e}")
            return self._generate_fallback_analysis(query, products_table)

    def _call_chat_completions(self, prompt: str) -> str:
        """调用chat/completions接口"""
        payload = {
            "model": "qwen2.5-vl-8b-instruct",
            "messages": [
                {
                    "role": "system",
                    "content": "你是一个专业的海关商品分类专家，擅长分析商品相似度和提供匹配建议。请用专业、清晰的语言回答。"
                },
                {
                    "role": "user",
                    "content": prompt
                }
            ],
            "max_tokens": 2000,
            "temperature": 0.3,
            "top_p": 0.9
        }

        response = requests.post(
            f"{self.qwen_api_url}/v1/chat/completions",
            json=payload,
            timeout=60
        )

        if response.status_code == 200:
            result = response.json()
            if 'choices' in result and len(result['choices']) > 0:
                return result['choices'][0]['message']['content']
            else:
                raise Exception("API返回格式异常")
        else:
            raise Exception(f"HTTP错误 {response.status_code}: {response.text}")

    def _generate_fallback_analysis(self, query: str, products_table: str) -> str:
        """生成备选分析报告"""
        return f"""
🔍 **商品匹配分析报告**

**查询商品:** {query}

**匹配结果概览:**
{products_table}

**简要分析:**
由于AI分析服务暂时不可用，以下是基于匹配结果的简要建议：

1. **查看相似度最高的商品** - 通常相似度 > 0.8 的商品是较好的匹配
2. **对比关键属性** - 重点关注规格、材质、用途的匹配程度
3. **检查HS编码** - 相同或相近的HS编码表示商品类别相似

**建议操作:**
- 如果找到高度匹配商品，可直接参考其海关分类
- 如匹配度一般，建议调整查询关键词或提供更多商品细节
- 可联系专业报关人员进行最终确认

*注: 此为基础分析，建议结合专业判断使用。*
"""

    def query_products(self, product_names: List[str], top_k: int = 5) -> Dict[str, Dict]:
        """查询产品 - 高效纯文本版本"""
        results = {}

        for query in product_names:
            print(f"🔍 处理查询: {query}")

            # 1. 检索阶段
            start_time = time.time()
            similar_products = self.retrieve_similar_products(query, top_k)
            retrieval_time = time.time() - start_time

            # 2. 生成阶段
            start_time = time.time()
            if similar_products:
                # 格式化产品表格
                products_table = self.format_products_table(similar_products)
                # 调用Qwen-VL进行分析
                response = self.call_qwen_vl_with_text(query, products_table)
            else:
                response = f"❌ 未找到与 '{query}' 相似的商品。"
                products_table = ""
            generation_time = time.time() - start_time

            results[query] = {
                'retrieved_products': similar_products,
                'products_table': products_table,
                'generated_response': response,
                'retrieval_time': retrieval_time,
                'generation_time': generation_time,
                'total_time': retrieval_time + generation_time
            }

            print(f"✅ 查询完成: 检索{retrieval_time:.2f}s + 生成{generation_time:.2f}s")

            # 短暂延迟避免API限制
            time.sleep(0.5)

        return results

    def print_detailed_results(self, results: Dict[str, Dict]):
        """打印详细结果"""
        for query, result in results.items():
            print(f"\n{'=' * 100}")
            print(f"📦 查询: {query}")
            print(f"{'=' * 100}")
            print(
                f"⏱️  总耗时: {result['total_time']:.2f}s (检索: {result['retrieval_time']:.2f}s, 生成: {result['generation_time']:.2f}s)")

            print(f"\n📊 检索结果表格:")
            print(result['products_table'])

            print(f"\n🤖 Qwen2.5-VL专业分析:")
            print(result['generated_response'])

            print(f"\n{'=' * 100}")


# 批量处理版本
class BatchQwenVLRAG(QwenVLRAGSystem):
    """批量处理优化的RAG系统"""

    def batch_query_products(self, product_names: List[str], top_k: int = 5, batch_size: int = 5) -> Dict[str, Dict]:
        """批量查询产品"""
        all_results = {}

        for i in range(0, len(product_names), batch_size):
            batch = product_names[i:i + batch_size]
            print(f"\n🔄 处理批次 {i // batch_size + 1}: {batch}")

            batch_results = self.query_products(batch, top_k)
            all_results.update(batch_results)

            # 保存批次结果
            self._save_batch_results(batch_results, i // batch_size + 1)

            # 批次间延迟
            if i + batch_size < len(product_names):
                print("⏳ 批次间休息...")
                time.sleep(2)

        return all_results

    def _save_batch_results(self, results: Dict[str, Dict], batch_num: int):
        """保存批次结果"""
        filename = f"batch_results_{batch_num}.json"
        try:
            with open(filename, 'w', encoding='utf-8') as f:
                # 转换无法序列化的对象
                serializable_results = {}
                for query, result in results.items():
                    serializable_results[query] = {
                        'retrieved_products': [
                            {
                                'product': {k: v for k, v in product.items()},
                                'score': score
                            }
                            for product, score in result['retrieved_products']
                        ],
                        'products_table': result['products_table'],
                        'generated_response': result['generated_response'],
                        'retrieval_time': result['retrieval_time'],
                        'generation_time': result['generation_time'],
                        'total_time': result['total_time']
                    }

                json.dump(serializable_results, f, ensure_ascii=False, indent=2)
            print(f"💾 批次 {batch_num} 结果已保存到 {filename}")
        except Exception as e:
            print(f"❌ 保存批次结果失败: {e}")


# 使用示例
def main():
    """主函数示例"""
    # 配置参数
    db_path = "path/to/your/清关资料.accdb"  # 替换为你的数据库路径
    qwen_api_url = "http://localhost:8000"  # 你的Qwen服务器地址

    print("🚀 初始化Qwen2.5-VL-8B RAG系统...")

    try:
        # 初始化RAG系统
        rag_system = QwenVLRAGSystem(db_path, qwen_api_url)

        # 测试查询
        test_queries = [
            "toy dog",
            "LED light bulb",
            "laptop computer",
            "plastic toy",
            "electric fan"
        ]

        print("\n🎯 开始商品匹配分析...")
        results = rag_system.query_products(test_queries, top_k=5)

        # 打印结果
        rag_system.print_detailed_results(results)

    except Exception as e:
        print(f"❌ 系统初始化失败: {e}")


def batch_processing_example():
    """批量处理示例"""
    db_path = "path/to/your/清关资料.accdb"
    qwen_api_url = "http://localhost:8000"

    try:
        rag_system = BatchQwenVLRAG(db_path, qwen_api_url)

        # 从文件读取查询列表或使用示例查询
        queries = [
            "toy dog", "LED bulb", "laptop", "plastic toy", "electric device",
            "children toy", "light fixture", "computer accessory", "household appliance"
        ]

        print(f"📋 开始批量处理 {len(queries)} 个查询...")
        results = rag_system.batch_query_products(queries, top_k=3, batch_size=3)

        print(f"\n✅ 批量处理完成，共处理 {len(results)} 个查询")

    except Exception as e:
        print(f"❌ 批量处理失败: {e}")


if __name__ == "__main__":
    # 选择运行模式
    print("请选择运行模式:")
    print("1. 单次查询示例")
    print("2. 批量处理示例")

    choice = input("请输入选择 (1-2): ").strip()

    if choice == "2":
        batch_processing_example()
    else:
        main()